#!/usr/bin/python2.6
# -*- coding: utf-8 -*-
# Copyright (c) 2008 Qtrac Ltd. All rights reserved.
# This program or module is free software: you can redistribute it and/or
# modify it under the terms of the GNU General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version. It is provided for educational
# purposes and is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# General Public License for more details.


import sys
import pickle
import string
import os
import io
import numpy as np
from scipy import stats
import time
import CaseConAlgo

def epiCaseControlOAInit():#TRASH??

	#Build jobs
	jobCC = {}
	jobCC["Project"]="To be inserted from the GUI"#TODO from GUI done
	
	pklfiles = ["epi_results1.pkl","epi_results6.pkl","epi_results10.pkl","epi_results14.pkl"]#TODO from GUI done
	epi_select = ["epi_results1.pkl","epi_results6.pkl"]#TODO from GUI done

	default_significance_level = 0.05#TODO from GUI ???
	jobCC["NomSignLevel"]=default_significance_level
	
	for epiSel in epi_select:
		epi_rem = pklfiles
		jobCC["ClassPrime"]=epiSel
		epi_rem.remove(jobCC["ClassPrime"])#.remove(epiSel), but from previous epi_rem file i.e. successive removal of pklfiles. Thats not what we want so we have append the epiSel-file
		jobCC["ClassRem"]=epi_rem
		
		#Listing traits TODO should pop up in GUI
		trait_list = []
		prime_res = open(jobCC["ClassPrime"],"rb")
		prime_data = pickle.load(prime_res)
		jobCC["PrimeData"]=prime_data
		for trait in prime_data["epistasis_data"]:#trait valg? GUI? 
			trait_select = trait["trait_name"]
			trait_list.append(trait_select)
		jobCC["TraitSel"]=trait_list

		epiCaseControlOATr(jobCC)
		
		#Reestablish class list
		epi_rem.append(jobCC["ClassPrime"])#Odd consequence of remove-statment



def start(primary_pkl_files, all_pkl_files, significance_level, trait_list, project_name, analysis_root_dir, callback=None):
		
		#jobCC["ClassPrime"]=epiSel
		#epi_rem.remove(jobCC["ClassPrime"])#.remove(epiSel), but from previous epi_rem file i.e. successive removal of pklfiles. Thats not what we want so we have append the epiSel-file
		#jobCC["ClassRem"]=epi_rem
		
		#Listing traits TODO should pop up in GUI
		#trait_list = []
		##prime_res = open(jobCC["ClassPrime"],"rb")
		#prime_data = pickle.load(prime_res)
		#jobCC["PrimeData"]=prime_data
		
	for i in range(len(primary_pkl_files)):
		pf = primary_pkl_files[i]
		jobCC = {}
		prime_res = open(pf)
		prime_data = pickle.load(prime_res)
		jobCC["PrimeData"]=prime_data
		all_other = list(all_pkl_files)
		all_other.remove(pf)
		jobCC["ClassRem"] = all_other
		
		jobCC["Project"]=project_name
		#jobCC["Pklfiles"] = class_pkl_files 
		jobCC["NomSignLevel"]= significance_level 
		jobCC["TraitSel"]=trait_list
		
		epiCaseControlOATr(jobCC, analysis_root_dir)
		if callback:
			callback(progress=int(float(i)/float(len(primary_pkl_files)))*100, message= "running %i of %i" % (i, len(primary_pkl_files)))
		#epi_rem.append(jobCC["ClassPrime"])#Odd consequence of remove-statment

		
	
	
def epiCaseControlOATr(jobCC, analysis_root_dir):

	"""For each Trait selected, classes are importet and trimmed, and all calculations are initiated"""

	for tr in jobCC["TraitSel"]:#TODO replace with GUI choises of traits, a job is defined!	
		new_summary_dir = os.path.join(analysis_root_dir, "Case-control One by All "+ tr +"/")
		if not os.path.isdir(new_summary_dir):#create output dir
			os.mkdir(new_summary_dir)	
		jobCC["OutputDir"]=new_summary_dir
		jobCC["trait"] = tr
		class_contentP = []#content of the selected file
		class_contentR = []#content of the remaining files		
		class_numbP = []
		class_nameP = []		
		
		#Content of selected class
#		prime_res = open(jobCC["ClassPrime"],"rb")
#		prime_data = pickle.load(prime_res)
		prime_data = jobCC["PrimeData"]	
		class_select = prime_data["class"]
		for trait in prime_data["epistasis_data"]:#trait valg? GUI? 
			trait_select = trait["trait_name"]
			if tr == trait_select:
				class_numbP.append(int(class_select))
				prime_fil = "GAfile" + str(class_select)
				jobCC["GAclass"]= prime_fil
				class_nameP.append(prime_fil)
				prime_fil = trait["AllGenFreq"]#A list of genotype and allel frequences as strings
				for ga in range(0,len(prime_fil)):
					al_li = prime_fil[ga].split()
					class_contentP.append(al_li)

		jobCC["ClassListNoP"]=class_numbP		
		jobCC["ClassListDataP"]=class_contentP



		#Content of the remaining files
		class_contentR = []		
		class_numbR = []
		class_nameR = []		
		class_list_numR = []
		for epi in jobCC["ClassRem"]:#Remaining pkl-files from GUI? No, automatic here
			rem_res = open(epi,"rb")
			rem_data = pickle.load(rem_res)
			class_select = rem_data["class"]
			for trait in rem_data["epistasis_data"]:#trait valg?  
				trait_select = trait["trait_name"]
				if tr == trait_select:
					class_numbR.append(int(class_select))

					rem_fil = "GAfile" + str(class_select)
					class_nameR.append(rem_fil)
					rem_fil = trait["AllGenFreq"]#A list of genotype and allel frequences as strings
					al_fil_mod = []
					for ga in range(0,len(rem_fil)):
						al_li = rem_fil[ga].split()
						al_fil_mod.append(al_li)
						class_contentR.append(al_li)#OK only used for sampling snps
					class_list_numR.append(al_fil_mod)#list of trimmed files
		class_numbR.sort()
		jobCC["ClassListNoR"]=class_numbR	
		jobCC["ClassListDataR"]=class_contentR
		jobCC["ClassListName"]=class_nameR
		jobCC["ClassListNumb"]=class_list_numR

		#Commen snps an significance levels
		case_control_snps(jobCC)

		#Create output files
		create_output_files(jobCC)

		#Start calculation
		cascon_select_data(jobCC)		

def case_control_snps(jobCC):

	"Find common snps"
	#Prime epi_results
	epifil_number1=jobCC["ClassListDataP"]
	jobCC["Raw1"] = epifil_number1
	snplist1 = CaseConAlgo.snp_list_create(epifil_number1)

	#Remaining epi_results
	epifil_number2=jobCC["ClassListDataR"]
	jobCC["Raw2"] = epifil_number2
	snplist2 = CaseConAlgo.snp_list_create(epifil_number2)

	#Only snps present both epifiles are of interest
	snp_list_common = list(set(snplist1[0]) & set(snplist2[0]))
	jobCC["SnpCommon"]=snp_list_common

	#Significance levels based on number of snps-tested
	sign_level_nominal= jobCC["NomSignLevel"]
	numb_tests = len(snp_list_common)*(len(snp_list_common)-1)/2
	if numb_tests>0:
		sign_level_bonf = sign_level_nominal/numb_tests
	else:
		sign_level_bonf = "No common SNPs"


	jobCC["SignNominal"]=sign_level_nominal	
	jobCC["SignBonf"]=sign_level_bonf


def create_output_files(jobCC):

	#Create output files
	#Output nominal
	
	print jobCC["ClassListNoP"]
	
	print jobCC["ClassListNoP"][0]
	print jobCC["ClassListNoR"]
	print jobCC["ClassListDataR"]
	print jobCC["ClassListName"]
	print jobCC["ClassListNumb"]

	
	snps_summary_00N = open(jobCC["OutputDir"]+"/"+"GAfile "+str(jobCC["ClassListNoP"][0]) + " vs All" +jobCC["trait"]+" Nominal", "w")
	jobCC["SumAllListN"]=snps_summary_00N	
	summ001 = time.asctime()+"\n\n"
	summ002 = "All SNPs in the selected file are compared with the all remaining files in the project"+"\n\n"
	summ002b = "Selected file : "+ jobCC["GAclass"] +"\n\n"
	summ002d = "Population files: "+str(jobCC["ClassListNoR"])+ "\n\n"
	summ002c = "Project: "+jobCC["Project"]+ "\n\n"
	summ002e = "Trait: "+jobCC["trait"]+ "\n\n"
	
	summ002a = "Number of SNPs compared: "+str(len(jobCC["SnpCommon"]))+ "\n\n"
	
	jobCC["SumAllListN"].write(summ001)
	jobCC["SumAllListN"].write(summ002)
	jobCC["SumAllListN"].write(summ002c)
	jobCC["SumAllListN"].write(summ002d)
	jobCC["SumAllListN"].write(summ002b)		
	jobCC["SumAllListN"].write(summ002e)			
	jobCC["SumAllListN"].write(summ002a)
	summ003 = "Significance level (Nominal): "+str(jobCC["SignNominal"])+"\n\n"
	summ004 = "{0} {1} {2} {3} {4} {5}\n".format("\t","\t"+jobCC["GAclass"]+"\t","\t","\t","\t","Population")
	summ004a = "{0} {1} {2} {3} {4} {5} {6} {7} {8} {9} {10}\n\n".format("SNP"+"\t","Gene"+"\t","AA"+"\t","Aa"+"\t","aa"+"\t","\t","BB"+"\t","Bb"+"\t","bb"+"\t","Stat_value"+"\t","p_value"+"\t")

	jobCC["SumAllListN"].write(summ003)
	jobCC["SumAllListN"].write(summ004)
	jobCC["SumAllListN"].write(summ004a)
	
	#Output Bonferoni
	snps_summary_00B = open(jobCC["OutputDir"]+"/"+"GAfiles "+str(jobCC["ClassListNoP"][0])+ " vs All"+" Bonferoni", "w")
	jobCC["SumAllListB"]=snps_summary_00B	
	summ001 = time.asctime()+"\n\n"
#	summ002 = "All SNPs compared pair-wise between subpopulations"+"\n\n"
	jobCC["SumAllListB"].write(summ001)
	jobCC["SumAllListB"].write(summ002)
	jobCC["SumAllListB"].write(summ002b)
	jobCC["SumAllListB"].write(summ002d)	
	jobCC["SumAllListB"].write(summ002c)	
	jobCC["SumAllListB"].write(summ002a)
	summ005 = "Significance level (Bonferoni): "+str(jobCC["SignBonf"])+"\n\n"
	jobCC["SumAllListB"].write(summ005)
	jobCC["SumAllListB"].write(summ004)
	jobCC["SumAllListB"].write(summ004a)
				
def cascon_select_data(jobCC):#own script

	#Create a file from sele selected file containing the selected snp
	for episnp in range(0,len(jobCC["SnpCommon"])):
		snp_select=jobCC["SnpCommon"][episnp]
		jobCC["SnpSelect"]=snp_select
		#Extract entries with selected snp from the selected file
		data_raw1_select=[]
		for epis in range(0,len(jobCC["Raw1"])):
			if jobCC["Raw1"][epis][0] == snp_select:
				data_raw1_select.append(jobCC["Raw1"][epis])
				genename1 = jobCC["Raw1"][epis][1]
			if jobCC["Raw1"][epis][7] == snp_select:
				data_raw1_select.append(jobCC["Raw1"][epis])
				genename1 = jobCC["Raw1"][epis][8]
		jobCC["DataCompile1"]=data_raw1_select
		jobCC["GeneName1"]=genename1

		#Extract entries with selected snp from the remaining files (population)????

		data_raw2_select=[]
		for epis in range(0,len(jobCC["Raw2"])):
			if jobCC["Raw2"][epis][0] == snp_select:
				data_raw2_select.append(jobCC["Raw2"][epis])
				genename2 = jobCC["Raw2"][epis][1]
			if jobCC["Raw2"][epis][7] == snp_select:
				data_raw2_select.append(jobCC["Raw2"][epis])
				genename2 = jobCC["Raw2"][epis][8]

		jobCC["DataCompile2"]=data_raw2_select
		
		#Call script to calulate case-control
		summ_list = cascon_select_calc(jobCC)

	
#def ggg():#temp out	
		#Results for snp_select to summary file, nominal significance
		if summ_list[1]<jobCC["SignNominal"]:
			snps_sign_head2 = "{0} {1} {2} {3} {4} {5} {6} {7} {8} {9} {10}\n".format(str(jobCC["SnpSelect"])+"\t",str(jobCC["GeneName1"])+"\t",str(jobCC["GenoList1"][0])+"\t",str(jobCC["GenoList1"][1])+"\t",str(jobCC["GenoList1"][2])+"\t","\t",str(jobCC["GenoList2"][0])+"\t",str(jobCC["GenoList2"][1])+"\t",str(jobCC["GenoList2"][2])+"\t",str(summ_list[0])+"\t",str(summ_list[1])+"\t")
			jobCC["SumAllListN"].write(snps_sign_head2)
			
		#Results for snp_select to summary file, Bonferoni corrected
		if summ_list[1] < jobCC["SignBonf"]:
			snps_sign_head2 = "{0} {1} {2} {3} {4} {5} {6} {7} {8} {9} {10}\n".format(str(jobCC["SnpSelect"])+"\t",str(jobCC["GeneName1"])+"\t",str(jobCC["GenoList1"][0])+"\t",str(jobCC["GenoList1"][1])+"\t",str(jobCC["GenoList1"][2])+"\t","\t",str(jobCC["GenoList2"][0])+"\t",str(jobCC["GenoList2"][1])+"\t",str(jobCC["GenoList2"][2])+"\t",str(summ_list[0])+"\t",str(summ_list[1])+"\t")
		
#			snps_sign_head3 = "{0} {1} {2} {3} {4} {5} {6} {7} {8} {9} {10}\n".format(jobCC["SnpSelect"],jobCC["GeneName1"],jobCC["GenoList1"][0],jobCC["GenoList1"][1],jobCC["GenoList1"][2],"\t",jobCC["GenoList2"][0],jobCC["GenoList2"][1],jobCC["GenoList2"][2],summ_list[0],summ_list[1])
			jobCC["SumAllListB"].write(snps_sign_head2)
		
def cascon_select_calc(jobCC):

	"""Calcualtions are performed using the largest numbers of subjects for each SNP. There may be some minor differences as the allelfrequency files are enumerated for two-SNP interactions when calculating two-SNP epistasis. However, these minor differences are not of interest; rather the larger the number of subjects the power of the case-control calculations will increase. This is justified, as lower numbers just mean that the genotype has beee filtered, but all filtered genotypes are actually present (and maybe more) in the basic suppopulation.
#Therfore consider using the basic populations, not those filtered by epistasis calculations"""

	#List genotypes from selected file. Necessary as the selected SNP may in column 0 or 7
	geno_types1 = []
	#From column 0..	
	for gt in range(0,len(jobCC["DataCompile1"])):
		if jobCC["DataCompile1"][gt][0] == jobCC["SnpSelect"]:
			geno_types1.append(jobCC["DataCompile1"][gt][2:5])#strange, only the first 3 to be used
	#From column 7..	
	for gt in range(0,len(jobCC["DataCompile1"])):
		if jobCC["DataCompile1"][gt][7] == jobCC["SnpSelect"]:
			geno_types1.append(jobCC["DataCompile1"][gt][9:12])
	#Find max number of subjects in the selected file
	genot1_list = []
	max_numb = 0	
	for gnt in range(0,len(geno_types1)):
		gg= geno_types1[gnt]
		gt_num = int(gg[0])+int(gg[1])+int(gg[2])		
		if gt_num>max_numb:
			max_numb = gt_num
			genot1_list = [int(gg[0]),int(gg[1]),int(gg[2])]		
	jobCC["GenoList1"]=genot1_list

	#List from remaining files, has to be treated individual as we are interesed in the sum genotypes/subjects
#	geno_types2 = []
	genot2_list = []#final list

	gg_list = []#cummulated list
	for ref in range(0,len(jobCC["ClassListNumb"])):
		geno_types2 = []
		for gt in range(0,len(jobCC["ClassListNumb"][ref])):
			#From column 0..	
			if jobCC["ClassListNumb"][ref][gt][0] == jobCC["SnpSelect"]:
				geno_types2.append(jobCC["ClassListNumb"][ref][gt][2:5])

			#From column 7..	
			if jobCC["ClassListNumb"][ref][gt][7] == str(jobCC["SnpSelect"]):
				geno_types2.append(jobCC["ClassListNumb"][ref][gt][9:12])

	#Find max number of subjects in the remaining files
		max_numb = 0		
		genot2_max=[0,0,0]
		for gnt in range(0,len(geno_types2)):
			gg= geno_types2[gnt]
			gt_num = int(gg[0])+int(gg[1])+int(gg[2])	
			if gt_num>max_numb:
				max_numb = gt_num
				genot2_max = [int(gg[0]),int(gg[1]),int(gg[2])]		
		gg_list.append(genot2_max)

	g2list = np.matrix(gg_list)
	genot2_list.append(sum(g2list)[0,0])
	genot2_list.append(sum(g2list)[0,1])
	genot2_list.append(sum(g2list)[0,2])
		
	jobCC["GenoList2"]=genot2_list

	calcepi=CaseConAlgo.cascon_calc_proper(jobCC)
	
	return calcepi

#TRASH
if __name__ == '__main__':
    start(primary_pkl_files=["epi_results1.pkl", "epi_results2.pkl"], all_pkl_files=["epi_results1.pkl", "epi_results2.pkl", "epi_results4.pkl", "epi_results3.pkl"], significance_level=0.05, trait_list=["ApoB"], project_name="test,", analysis_root_dir=".")

#epiCaseControlOAInit()


